import warnings
warnings.filterwarnings('ignore', 'Optimal rotation is not uniquely or poorly defined for the given sets of vectors.')

from tqdm import tqdm
import numpy as np
from scipy.spatial.transform import Rotation as R

from .utils import timer, rnd
from .typing import *


# 定义函数生成一定数量的点
def gen_points(x_range:Range, y_range:Range, z_range:Range, interval:float=0.2) -> Optional[Points]:
    # 计算点阵大小
    size_x = int((x_range[1] - x_range[0]) / interval) + 1 
    size_y = int((y_range[1] - y_range[0]) / interval) + 1
    size_z = int((z_range[1] - z_range[0]) / interval) + 1
    sizes = [size_x, size_y, size_z]
    if 1 in sizes: return

    # 构建坐标点阵：[[0, 0, 0], [0, 0, 1], ..., [nX-1, xY-1, xZ-1]]
    zeros = np.zeros(sizes)                         # [nX, xY, xZ]
    grid = np.transpose(np.where(zeros == 0.0))     # [nX*xY*xZ, D=3]
    grid = grid.astype(DTYPE)

    # 重映射数据： coord_rang: [0, nAxis] => v_range: (vmin, vmax)
    def map_data(idx:int, size:int, vrng:Range):
        grid[:, idx] = (grid[:, idx] / (size - 1)) * (vrng[1] - vrng[0]) + vrng[0]

    map_data(0, size_x, x_range)
    map_data(1, size_y, y_range)
    map_data(2, size_z, z_range)

    return grid


def count_points(atom_contacts:Points, atom_centers:Dict[Tuple[str, float], Point]) -> ndarray:
    # 计算所有点与原子中心之间的距离
    atom_info = np.asarray(list(atom_centers.values()))                # 原子的坐标
    radii = np.asarray([radius for _, radius in atom_centers.keys()])  # 原子的半径
    d = atom_info[:, np.newaxis] - atom_contacts
    dists = np.linalg.norm(d, axis=2)               # 这一步较慢，但似乎是唯一的解决方案

    # 找出距离小于等于r的点
    within_r = dists <= radii[:, np.newaxis]
    
    # 计算每个点在多少个原子内部。根本不需要原子坐标了
    atom_counts = np.sum(within_r, axis=0)
    
    return atom_counts


def determine_inner(a_pair:Tuple[str, str], a_pair_xyz:Dict[str, Point], a_pair_info:Dict[str, float], interval:float=0.2) -> Tuple[dict, dict]:
    volume_pair = {}
    xyz = {}
    atom1, atom2 = a_pair
    
    # 把原子1放在坐标原点
    for atom in a_pair_xyz:
        xyz[atom] = a_pair_xyz[atom] - a_pair_xyz[atom1]
    
    # 原子2到原子1的向量
    vector = xyz[atom2]
    # 计算这个向量与X轴之间的旋转
    rotation, _ = R.align_vectors([[1, 0, 0]], [vector])
    # 旋转所有点，这样原子2会落到X轴上，可能在正方向也可能在负方向
    centers_array = np.asarray(list(xyz.values()))
    rotated_centers_array = rotation.apply(centers_array)
    # 旋转后的点的坐标
    centers = {key: value for key, value in zip(xyz.keys(), rotated_centers_array)}
    
    R1 = DTYPE(a_pair_info[atom1])
    R2 = DTYPE(a_pair_info[atom2])
    S  = DTYPE(centers[atom2][0])    # 两个原子之间的距离（但可能是负值，取决于旋转方向）
    
    # 确定点阵的范围
    if S > 0:
        x_max = 0 + R1
        x_min = S - R2
    elif S < 0:
        x_max = S + R2
        x_min = 0 - R1
    
    if R1 >= R2:
        y_max = + R1
        y_min = - R1
        z_max = + R1
        z_min = - R1
    elif R2 > R1:
        y_max = + R2
        y_min = - R2
        z_max = + R2
        z_min = - R2

    # 计算点阵的体积和点数，+0.1是为了避免太靠近边缘的点不精确
    eps = 0.1
    cube_volu: DTYPE = rnd((x_max - x_min + eps) * (y_max - y_min + eps) * (z_max - z_min + eps), 9)

    # 指定每个维度的范围
    x_range = (rnd(x_min - eps / 2), rnd(x_max + eps / 2))  
    y_range = (rnd(y_min - eps / 2), rnd(y_max + eps / 2))
    z_range = (rnd(z_min - eps / 2), rnd(z_max + eps / 2))

    # 点阵坐标
    points: Points = gen_points(x_range, y_range, z_range, interval=interval)
    if points is None: return {}
    
    # 原子1和原子2的坐标
    two_center = np.asarray([centers[atom1], centers[atom2]])
    
    # 点阵到原子1和原子2的距离
    dists = np.linalg.norm(points[:, np.newaxis] - two_center, axis=2)
    
    # 在原子1和原子2之内的点
    in_contact = points[(dists[:, 0] <= R1) & (dists[:, 1] <= R2)]
    atom1_centers = {(c, a_pair_info[c]): centers[c] for c in centers} ; del atom1_centers[(atom1, a_pair_info[atom1])]
    atom1_counts = count_points(in_contact, atom1_centers)

    # 每个点代表的体积
    a_point = cube_volu / len(points)
    
    # 最后计算接触体积
    volu: DTYPE = (1 / atom1_counts).sum()
    volume = rnd(volu * a_point)
    
    # 把反向的原子对也加进去，方便后期处理
    volume_pair[a_pair]       = volume
    volume_pair[a_pair[::-1]] = volume
    
    return volume_pair


def pdb_dotarray_volume(atom_df:DataFrame, interval:float=0.2, disable_print=False):
    @timer(disable_print=disable_print)
    def count_volume():
        # 原子名称和半径信息
        atoms: Series = atom_df['Residue'] + '_' + atom_df['Atom']
        radii = np.asarray(atom_df['R'])
        
        # 坐标的数组
        atom_coords: Points = np.stack([
            atom_df['x'].array,    # [N], DTYPE
            atom_df['y'].array,
            atom_df['z'].array,
        ], axis=-1)                 # [N, D=3], DTYPE
        
        atom_df['Name'] = atoms
        contacts_center = {i: j for i, j in zip(atoms, atom_coords)}
        contacts_dict = atom_df.set_index('Name')['R'].to_dict()
        
        # 计算所有原子之间的距离
        dists = np.linalg.norm(atom_coords[:, np.newaxis] - atom_coords, axis=2)

        # 创建一个字典，键是原子名称，值是与该原子接触的所有原子名称的列表
        # >0.00001是为了去掉该原子自己，因为自己到自己的距离是0
        eps = 0.00001
        contact_map = {
            atom1: atoms[(dists[i] > eps) & (dists[i] < radii[i] + radii)]
                for i, (atom1, radius) in enumerate(zip(atoms, radii))
        }
        
        contact_dict = {} # 构建成原子对
        for key, values in contact_map.items():
            key_residue = key.split('_')[0]
            for value in values:
                value_residue = value.split('_')[0]
                if key_residue != value_residue:
                    new_key = (key, value)
                    new_values = list(values) + [key]
                    contact_dict[new_key] = new_values
        
        Pairs_Volume = []
        for a_pair in tqdm(contact_dict, disable=disable_print):
            # 一个原子对，以及与它们接触的所有原子
            a_contact_dict = contact_dict[a_pair]
            a_pair_xyz  = {name: contacts_center[name] for name in a_contact_dict if name in contacts_center}   # 坐标
            a_pair_info = {name: contacts_dict  [name] for name in a_contact_dict if name in contacts_dict  }   # 半径、体积
            
            volume_pair = determine_inner(a_pair, a_pair_xyz, a_pair_info, interval)
            Pairs_Volume.append(volume_pair)
        
        volume = {} # 使用update快速合并字典
        for pair in Pairs_Volume:
            volume.update(pair)  

        return contact_dict, volume
    
    return count_volume()
